import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 硬件配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# 数据预处理（改进归一化参数）
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomRotation(10),  # 新增数据增强
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 调整后的归一化参数
])

# 数据集加载
train_dataset = torchvision.datasets.MNIST(root='./data', train=True,
                                          transform=transform, download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False,
                                         transform=transform)

# 数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


# 新增激活函数
class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super().__init__()
        self.inplace = inplace

    def forward(self, x):
        return x * nn.functional.relu6(x + 3, inplace=self.inplace) / 6


class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super().__init__()
        self.inplace = inplace

    def forward(self, x):
        return nn.functional.relu6(x + 3, inplace=self.inplace) / 6


# 新增SE模块
class SELayer(nn.Module):
    def __init__(self, channel, reduction=4):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel),
            h_sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


# 核心构建块
class InvertedResidual(nn.Module):
    def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
        super().__init__()
        assert stride in [1, 2]

        self.identity = stride == 1 and inp == oup

        layers = []
        # 扩展层
        if inp != hidden_dim:
            layers.append(nn.Conv2d(inp, hidden_dim, 1, bias=False))
            layers.append(nn.BatchNorm2d(hidden_dim))
            layers.append(h_swish() if use_hs else nn.ReLU(inplace=True))

        # 深度卷积
        layers.append(nn.Conv2d(
            hidden_dim, hidden_dim, kernel_size, stride,
            (kernel_size - 1) // 2, groups=hidden_dim, bias=False))
        layers.append(nn.BatchNorm2d(hidden_dim))
        layers.append(h_swish() if use_hs else nn.ReLU(inplace=True))

        # SE模块
        if use_se:
            layers.append(SELayer(hidden_dim))

        # 投影层
        layers.append(nn.Conv2d(hidden_dim, oup, 1, bias=False))
        layers.append(nn.BatchNorm2d(oup))

        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.identity:
            return x + self.conv(x)
        else:
            return self.conv(x)


# MobileNetV3-small实现
class MobileNetV3_Small(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # 初始层
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(16),
            h_swish(),

            # bneck配置 (expansion, out, kernel_size, stride, SE, HS)
            self._make_layer(16, 16, 16, 3, 2, False, False),
            self._make_layer(16, 72, 24, 3, 2, False, False),
            self._make_layer(24, 88, 24, 3, 1, False, False),
            self._make_layer(24, 96, 40, 5, 2, True, True),
            self._make_layer(40, 240, 40, 5, 1, True, True),
            self._make_layer(40, 240, 40, 5, 1, True, True),
            self._make_layer(40, 120, 48, 5, 1, True, True),
            self._make_layer(48, 144, 48, 5, 1, True, True),
            self._make_layer(48, 288, 96, 5, 2, True, True),
            self._make_layer(96, 576, 96, 5, 1, True, True),
            self._make_layer(96, 576, 96, 5, 1, True, True),
        )

        # 最后层
        self.conv = nn.Sequential(
            nn.Conv2d(96, 576, 1, stride=1, bias=False),
            nn.BatchNorm2d(576),
            h_swish()
        )
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Linear(576, 1024),
            h_swish(),
            nn.Dropout(0.2),
            nn.Linear(1024, num_classes)
        )

    def _make_layer(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
        return InvertedResidual(
            inp=inp,
            hidden_dim=hidden_dim,
            oup=oup,
            kernel_size=kernel_size,
            stride=stride,
            use_se=use_se,
            use_hs=use_hs
        )

    def forward(self, x):
        x = self.features(x)
        x = self.conv(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


# 初始化模型
model = MobileNetV3_Small().to(device)

# 初始化模型
print("开始训练")
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)  # 学习率调度

print("开始训练2")
# 训练循环
total_step = len(train_loader)
for epoch in range(num_epochs):
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 1 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Loss: {loss.item():.4f}')
    scheduler.step()

# 测试阶段
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print(f'测试准确率: {100 * correct / total:.2f}%')
